# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo

import enum
import hashlib
import hmac
import json
import logging
import time
import typing

import tornado.concurrent
import tornado.gen
import tornado.ioloop
import tornado.websocket

from coin.base.timestamp import get_timestamp
from coin.exchange.base.kr_rest.auth_base import AuthKey
from coin.exchange.base.rest.task_scheduler import Scheduler

HITBTC_WSS_URL = 'wss://api.hitbtc.com/api/2/ws'


class HitbtcOrderStatus(enum.Enum):
  NEW = 'new'
  SUSPENDED = 'suspended'
  PARTIALLY_FILLED = 'partiallyFilled'
  FILLED = 'filled'
  CANCELED = 'canceled'
  EXPIRED = 'expired'


class HitbtcReportType(enum.Enum):
  STATUS = 'status'
  NEW = 'new'
  CANCELED = 'canceled'
  EXPIRED = 'expired'
  SUSPENDED = 'suspended'
  TRADE = 'trade'
  REPLACED = 'replaced'


class HitbtcOrderSide(enum.Enum):
  BUY = 'buy'
  SELL = 'sell'


class HitbtcOrderType(enum.Enum):
  LIMIT = 'limit'
  MARKET = 'market'
  STOP_LIMIT = 'stopLimit'
  STOP_MARKET = 'stopMarket'


class HitbtcTimeInForce(enum.Enum):
  GTC = 'GTC'
  IOC = 'IOC'
  FOK = 'FOK'
  DAY = 'DAY'
  GTD = 'GTD'


class HitbtcWsPrivateClient(object):
  def __init__(
      self,
      key_file,
      ioloop=None,
      notification_callback=None,
      close_callback=None,
      on_ready=None,
      timeout=20,
      api_url=HITBTC_WSS_URL,
  ):
    self._key = AuthKey.from_file(key_file)

    self._ioloop = ioloop or tornado.ioloop.IOLoop.current()
    self._ws = None
    self._ready = False
    self._num_timeouts = 0

    self._pending_futures = {}
    self._request_id = int(time.time())

    self._notification_callback = notification_callback
    self._close_callback = close_callback or (lambda: None)
    self._on_ready = on_ready or (lambda: None)
    self._scheduler = Scheduler()
    self._api_url = api_url

    self._last_received_message_timestamp = None
    self._timeout = timeout * 1e9

    self._stop = False
    self._init_params = (
        key_file,
        ioloop,
        notification_callback,
        close_callback,
        on_ready,
        timeout,
        api_url,
    )
    self._logger = logging.getLogger(__name__)

  @property
  def num_timeouts(self):
    return self._num_timeouts

  @property
  def last_received_message_timestamp(self):
    return self._last_received_message_timestamp

  def is_ready(self):
    return self._ready

  def _check_futures(self):
    timestamp = get_timestamp()
    if self._stop:
      return

    for req_id, fut in self._pending_futures.copy().items():
      if timestamp - fut.timestamp > self._timeout:
        self._num_timeouts += 1
        fut = self._pending_futures.pop(req_id)
        fut.set_exception(TimeoutError('Websocket request timeout! (%s, %s)' %
                                       (req_id, fut.method)))

  def _get_next_request_id(self):
    self._request_id += 1
    res = self._request_id
    return res

  def start(self):
    self._scheduler.start()
    self._ioloop.add_callback(self._connect)

  def _write_message(self, msg):
    try:
      msg_str = json.dumps(msg)
      # self._logger.info('write message: %s', msg_str)
      self._ws.write_message(msg_str)
    except tornado.websocket.WebSocketClosedError as e:
      self._logger.error('Write message error, ws is close!')
      if self._pending_future is not None:
        self._pending_future.set_exception(e)
        self._pending_future = None

  @tornado.gen.coroutine
  def _connect(self):
    # noinspection PyBroadException
    try:
      self._logger.info('%s try to connect to the websocket.', type(self).__name__)
      self._ws = yield tornado.websocket.websocket_connect(self._api_url)
      self._logger.info('websocket connected!')
      self._ioloop.add_callback(self._run)
      yield self.login(basic=False)
      yield tornado.gen.sleep(5)
      yield self.subscribe_to_reports()
      self._ready = True
      self._on_ready()
    except Exception:
      self._logger.exception('%s connect to websocket failed!', type(self).__name__)
      self.stop()

  @tornado.gen.coroutine
  def _run(self):
    while not self._stop:
      msg = None
      # noinspection PyBroadException
      try:
        msg = yield self._ws.read_message()
      except Exception:
        self._logger.exception('Error reading ws message')
        self.stop()

      if msg is None:
        self._logger.error('Websocket is closed!')
        self.stop()
        return

      if len(msg) == 0:
        self._logger.info('Empty message received!')
        yield tornado.gen.sleep(0.01)
        continue

      # noinspection PyBroadException
      try:
        self.on_message(msg)
      except Exception:
        self._logger.exception('Exception raised when process message!')

  def stop(self):
    if self._stop:
      return

    self._stop = True
    self._ready = False
    self._ws.close()
    self._ws = None
    self._close_callback()

  def clone(self):
    return type(self)(*self._init_params)

  def on_message(self, raw_msg: bytes):
    self._last_received_message_timestamp = get_timestamp()
    if self._stop:
      return

    try:
      data = json.loads(raw_msg)
    except ValueError:
      self._logger.exception('Invalid json string!, %s', raw_msg)
      return

    if 'method' in data:  # notification
      self._notification_callback(data)
    elif 'result' in data or 'error' in data:  # request response
      request_id = data.get('id', None)
      if request_id is None:  # subscription
        if 'error' in data:
          self._logger.info('subscribe reports error!')
        else:
          self._logger.info('subscribe reports done!')
        return
      fut = self._pending_futures.pop(request_id, None)
      if fut is None:
        self._logger.error('Weird! request_id not found!')
      else:
        fut.set_result(data)
    else:
      self._logger.error('Skip invalid message: %s', data)

  def send_request(self, method, params=None, is_notification=False):
    if self._ws is None:
      self._logger.error('websocket not connected yet! method=%s, params=%s', method, params)
      return

    params = params or {}
    params = {key: value for key, value in params.items() if value is not None}
    request_object = {
        'jsonrpc': '2.0',
        'method': method,
        'params': params,
    }

    fut = None
    if not is_notification:
      request_id = self._get_next_request_id()
      request_object['id'] = request_id
      fut = tornado.concurrent.Future()
      fut.timestamp = get_timestamp()
      fut.method = method
      self._pending_futures[request_id] = fut
    self._ioloop.add_callback(self._write_message, request_object)
    return fut

  def login(self, basic=True):
    method = 'login'
    if basic:
      params = {
          'algo': 'BASIC',
          'pKey': self._key.access_key,
          'sKey': self._key.secret_key,
      }
    else:
      nonce = str(round(time.time() * 1000))
      signature = hmac.new(self._key.secret_key_bytes, nonce.encode(), hashlib.sha256).hexdigest()
      params = {
          'algo': 'HS256',
          'pKey': self._key.access_key,
          'nonce': nonce,
          'signature': signature,
      }
    return self.send_request(method, params)

  def subscribe_to_reports(self):
    method = 'subscribeReports'
    return self.send_request(method, is_notification=True)

  def place_new_order(
      self,
      *,
      symbol,
      price,
      quantity,
      order_side: HitbtcOrderSide,
      order_type: HitbtcOrderType,
      client_order_id,
      time_in_force: typing.Union[HitbtcTimeInForce, None] = None,
      stop_price=None,
      expire_time=None,
      strict_validate=True,
      post_only=False,
  ):
    method = 'newOrder'

    if time_in_force is not None:
      time_in_force = time_in_force.value

    params = {
        'clientOrderId': client_order_id,
        'symbol': symbol,
        'side': order_side.value,
        'type': order_type.value,
        'timeInForce': time_in_force,
        'quantity': quantity,
        'price': price,
        'stopPrice': stop_price,
        'expireTime': expire_time,
        'strictValidate': strict_validate,
        'postOnly': post_only,
    }
    return self.send_request(method, params)

  def cancel_order(self, client_order_id):
    method = 'cancelOrder'
    params = {
        'clientOrderId': client_order_id,
    }
    return self.send_request(method, params)

  def amend_order(self, client_order_id, request_client_order_id, quantity, price, strict_validate):
    method = 'cancelReplaceOrder'
    params = {
        'clientOrderId': client_order_id,
        'requestClientId': request_client_order_id,
        'quantity': quantity,
        'price': price,
        'strictValidate': strict_validate,
    }
    return self.send_request(method, params)

  def get_active_orders(self):
    method = 'getOrders'
    return self.send_request(method)

  def get_trading_balance(self):
    method = 'getTradingBalance'
    return self.send_request(method)
